package com.mrzhou.study.datasource;

import com.mrzhou.core.tool.utils.Func;
import lombok.SneakyThrows;
import org.springframework.beans.MutablePropertyValues;
import org.springframework.beans.factory.support.BeanDefinitionRegistry;
import org.springframework.beans.factory.support.GenericBeanDefinition;
import org.springframework.boot.jdbc.DataSourceBuilder;
import org.springframework.context.EnvironmentAware;
import org.springframework.context.annotation.Configuration;
import org.springframework.context.annotation.ImportBeanDefinitionRegistrar;
import org.springframework.core.env.Environment;
import org.springframework.core.type.AnnotationMetadata;

import javax.sql.DataSource;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

/**
 * 注册动态数据源并提供切换数据源的工具方法
 */
public class DynamicDataSourceConfiguration implements ImportBeanDefinitionRegistrar, EnvironmentAware {

    // 使用map结构保存数据源, ConcurrentHashMap是线程安全的
    private Map<String, Object> saveDataSource = new ConcurrentHashMap<>();
    // 指定默认数据源类型, 在配置文件没有配置数据源类型时,默认使用的SpringBoot自带数据源
    private final String DEFAULT_DATASOURCE_TYPE = "com.zaxxer.hikari.HikariDataSource";
    /** 项目默认读取的数据源前缀 */
    private static List<String> DEFAULT_DATASOURCE_PREFIX = new ArrayList<>();
    private static String DEFAULT_DATASOURCE_ID ;

    @SneakyThrows
    @Override
    public void setEnvironment(Environment environment) {
        initDataSource(environment);
    }

    @Override
    public void registerBeanDefinitions(AnnotationMetadata annotationMetadata,
                                        BeanDefinitionRegistry beanDefinitionRegistry) {
        Map<String, Object> targetDataSources = new HashMap<>();
        targetDataSources.putAll(saveDataSource);
        for (String key : saveDataSource.keySet()) {
            DynamicDataSourceContextHolder.dataSourceIds.add(key);
        }

        // 创建DynamicDataSource
        GenericBeanDefinition beanDefinition = new GenericBeanDefinition();
        beanDefinition.setBeanClass(DynamicDataSource.class);
        beanDefinition.setSynthetic(true);
        MutablePropertyValues mpv = beanDefinition.getPropertyValues();
        // 默认数据源
        mpv.addPropertyValue("defaultTargetDataSource", targetDataSources.get(DEFAULT_DATASOURCE_ID));
        mpv.addPropertyValue("targetDataSources", targetDataSources);
        // 注册 - BeanDefinitionRegistry
        beanDefinitionRegistry.registerBeanDefinition("dataSource", beanDefinition);
    }

    private void initDataSource(Environment environment) throws Exception {
        // 从配置文件中读取数据源的配置
        String namesStr = environment.getProperty("spring.datasource.names");
        String confPrefix = "spring.datasource.";
        List<String> names = handlerPrefix(namesStr);
        DEFAULT_DATASOURCE_ID = names.get(0); // 默认使用第一个作为默认的数据源
        for(String name: names) {
            String prefix = confPrefix + name + ".";
            // 保存多数据源配置
            Map<String, Object> dataSourceMap = new HashMap<>();
            dataSourceMap.put("type", environment.getProperty(prefix + "type"));
            dataSourceMap.put("driver", environment.getProperty(prefix + "driver-class-name"));
            dataSourceMap.put("url", environment.getProperty(prefix + "url"));
            dataSourceMap.put("username", environment.getProperty(prefix + "username"));
            dataSourceMap.put("password", environment.getProperty(prefix + "password"));
            DataSource dataSource = builderDataSource(dataSourceMap);
            saveDataSource.put(name, dataSource);
        }
    }

    /**
     * 构建数据源
     */
    private DataSource builderDataSource(Map<String, Object> dataSourceConf) throws Exception {
        Object type = dataSourceConf.get("type");
        if(type == null) {
            type =DEFAULT_DATASOURCE_TYPE;
        }
        Class<? extends DataSource> dataSourceType = (Class<? extends DataSource>) Class.forName((String) type);
        String driverClassName = dataSourceConf.get("driver").toString();
        String url = dataSourceConf.get("url").toString();
        String username = dataSourceConf.get("username").toString();
        String password = dataSourceConf.get("password").toString();
        // 自定义DataSource配置
        DataSourceBuilder<?> builder = DataSourceBuilder.create()
                .type(dataSourceType).driverClassName(driverClassName)
                .url(url).username(username).password(password);
        return builder.build();
    }

    private List<String> handlerPrefix(String prefix) {
        List<String> names = Func.toArray(prefix);
        if(names.isEmpty()) {
            names = DEFAULT_DATASOURCE_PREFIX;
        }
        return names;
    }

    static {
        DEFAULT_DATASOURCE_ID = "master";
        DEFAULT_DATASOURCE_PREFIX.add("master");
        DEFAULT_DATASOURCE_PREFIX.add("slave");
    }

}
